# %%
from datetime import datetime, timezone
from itertools import batched
import time
from loguru import logger
from pydantic import BaseModel, Field
from pymongo import UpdateOne
from tqdm import tqdm
from conn import homologous_clip, mysql_engine, clip_col, clip_annotated_record
from sqlalchemy import text


class HomologousClipBase(BaseModel):
    clip_id: str = Field(
        title="Clip ID",
    )
    annotate_type_list: list[int] = Field(
        default=[],
        title="标注类型列表",
    )
    curated_annotate_type_list: list[int] = Field(
        default=[],
        title="送标标注类型列表",
    )
    labeled_annotate_type_list: list[int] = Field(
        default=[],
        title="回标标注类型列表",
    )
    confirmed_annotate_type_list: list[int] = Field(
        default=[],
        title="回标确认标注类型列表",
    )
    name: str | None = Field(
        default=None,
        title="Clip Name",
    )
    tags: dict[str, list[str | int]] | None = Field(
        default=None,
        title="标签",
    )
    collect_at: str | None = Field(
        default=None,
        title="采集日期",
    )
    collect_time: int | None = Field(
        default=None,
        title="采集时间",
    )


CHAR_ARRAY = [
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
    "a",
    "b",
    "c",
    "d",
    "e",
    "f",
    "g",
]

# %%
# 查询所有送标的 clip
BATCH = 1000
clip_map = {}
with mysql_engine.connect() as conn:
    for i in range(len(CHAR_ARRAY) - 1):
        start_char, end_char = CHAR_ARRAY[i], CHAR_ARRAY[i + 1]
        sql = text(
            """
            select clip_id, annotate_type
            from send_label_record_tab 
            where clip_id >= :start_char and clip_id < :end_char 
            group by clip_id, annotate_type
            """
        )
        sql = sql.bindparams(start_char=start_char, end_char=end_char)
        t1 = time.perf_counter()
        rows = conn.execute(sql).fetchall()
        t2 = time.perf_counter()
        logger.info(f"[{start_char}, {end_char}) cost {t2 - t1:.3f}s, {len(rows)}")
        if len(rows) == 0:
            break
        for clip_id, annotate_type in rows:
            if clip_id not in clip_map:
                clip_map[clip_id] = HomologousClipBase(clip_id=clip_id)
            clip_map[clip_id].curated_annotate_type_list.append(annotate_type)


clip_ids = list(clip_map.keys())
logger.info(f"{len(clip_ids)} clips")

# %%
# 补全 clip 的其他信息

for batch_clip_ids in tqdm(batched(clip_ids, BATCH), total=len(clip_ids) // BATCH):
    batch_clip_infos = clip_col.find(
        {"_id": {"$in": batch_clip_ids}},
        {"name": 1, "tags": 1, "collect_at": 1, "collect_time": 1},
    )
    for clip_info in batch_clip_infos:
        clip_id = clip_info["_id"]
        clip = clip_map[clip_id]
        clip.name = clip_info["name"]
        clip.tags = clip_info.get("tags", None)
        clip.collect_at = clip_info["collect_at"]
        clip.collect_time = clip_info["collect_time"]

# %%
# 查询所有回标的 clip
for batch_clip_ids in tqdm(batched(clip_ids, BATCH), total=len(clip_ids) // BATCH):
    batch_clip_infos = clip_annotated_record.find(
        {"clip_id": {"$in": batch_clip_ids}},
        {"state": 1, "clip_id": 1, "annotate_type": 1},
    )
    for clip_info in batch_clip_infos:
        clip_id = clip_info["clip_id"]
        clip = clip_map[clip_id]
        if clip_info["annotate_type"] not in clip.labeled_annotate_type_list:
            clip.labeled_annotate_type_list.append(clip_info["annotate_type"])
        if (
            clip_info["state"] == 2
            and clip_info["annotate_type"] not in clip.confirmed_annotate_type_list
        ):
            clip.confirmed_annotate_type_list.append(clip_info["annotate_type"])


# 看下数据情况
for clip_id in clip_ids[:10]:
    logger.info(clip_map[clip_id].model_dump())

# %%
# 准备插入和更新
for batch_clip_ids in tqdm(batched(clip_ids, BATCH), total=len(clip_ids) // BATCH):
    clip_list = [
        clip_map[clip_id].model_dump(exclude=("annotate_type_list",))
        for clip_id in batch_clip_ids
    ]
    ops = [
        UpdateOne(
            {"clip_id": clip["clip_id"]},  # 查询条件，根据 clip_id 匹配文档
            {
                "$setOnInsert": {"created_at": datetime.now(tz=timezone.utc)},
                "$set": clip,
            },
            upsert=True,  # 启用 upsert
        )
        for clip in clip_list
    ]
    homologous_clip.bulk_write(ops, ordered=False)
